# -*- coding=utf-8 -*-
#!/usr/bin/env python

#从多个文件中查找指定的位点
import sys
import os
import re
import json
import argparse
from collections import defaultdict

from utils import TitleParser

def reading_file_list(afile):
    file_list = []
    with open(afile,'r') as indata:
        for line in indata:
            file_list.append(line.strip())
    return file_list

def init_index(file_of_list,idx_name,by):
    """
    reading in files with Chr POS Ref Alt
    and generate a index json file
    """
    file_list = reading_file_list(file_of_list)
    idx_dict = defaultdict(list)
    for f in file_list:
        re_str = re.search('converted.([\w]+)?_(indel|snp).xls',os.path.basename(f))
        if re_str is not None:
            tag = re_str.group(1)
        else:
            raise Exception('Unknown tag {tag}'.format(tag=f))

        with open(f,'r') as indata:
            title = indata.readline()
            tp = TitleParser(title)
            for line in indata:
                line_list = line.strip().split('\t')
                chrom = tp.get_field(line_list,'Chrom',check=False)
                pos = tp.get_field(line_list,'POS',check=False)
                ref = tp.get_field(line_list,'Ref',check=False)
                alt = tp.get_field(line_list,'Alt',check=False)
                genes = tp.get_field(line_list,'Gene.refGene',check=False).strip('"').split(",")
                # print(genes)
                if by == "site":
                    try:
                        key = "_".join([chrom,pos,ref,alt])
                    except:
                        print([chrom,pos,ref,alt])
                        exit(0)
                    if not idx_dict.get(key):
                        idx_dict[key].append(tag)
                    elif tag not in idx_dict.get(key):
                        idx_dict[key].append(tag)

                elif by == "gene":
                    for gene in genes:
                        key = gene
                        # print(key)
                        if not idx_dict.get(key):
                            idx_dict[key].append(tag)
                        elif tag not in idx_dict.get(key):
                            idx_dict[key].append(tag)

    with open(idx_name,'w') as odata:
        json.dump(idx_dict,odata,indent=2,ensure_ascii=False)
    print("init index output to {ofile}".format(ofile=idx_name))

def add_tags(infile,ofile,idx_name,by):
    """
    add tag to chrom_pos_ref_alt key
    """
    with open(idx_name,'r') as idx_data,\
        open(infile,'r') as indata,\
        open(ofile,'w') as odata:
        idx_dict = json.load(idx_data)
        title = indata.readline()
        tp = TitleParser(title)
        otitle = title.strip() + "\tShare_num_in_WES\tShare_samples\n"
        odata.write(otitle)
        for line in indata:
            line_list = line.strip().split('\t')
            chrom = tp.get_field(line_list,'Chrom',check=False)
            pos = tp.get_field(line_list,'POS',check=False)
            ref = tp.get_field(line_list,'Ref',check=False)
            alt = tp.get_field(line_list,'Alt',check=False)
            genes = tp.get_field(line_list,'GeneName',check=False).split(",")
            if by == "site":
                keys = ["_".join([chrom,pos,ref,alt])]
            elif by == "gene":
                keys = genes
            
            
            final_tag = set()
            for key in keys:
                tag = idx_dict.get(key)
                if tag:
                    for t in tag:
                        final_tag.add(t) 
            if not final_tag:
                final_tag_num = 0
                final_tag = "."
            else:
                final_tag_num = len(final_tag)
                final_tag = ",".join(final_tag)
            oline_list = line_list + [str(final_tag_num),final_tag]
            # print(oline_list)
            # exit(0)
            odata.write("\t".join(oline_list) + "\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""
        Usage:
        #init index file
        python sites_search.py --init file_list --ofile index_file
        #add  WES samples tags
        python sites_search.py --add_tag idx_file --infile infile --ofile ofile
        """
    )
    parser.add_argument("--init","-ini",help="a init file list")
    parser.add_argument("--add_tag","-a",help="a init index file used to add tag")
    parser.add_argument("--infile","-i",help="a file needed to add tags")
    parser.add_argument("--ofile","-o",help="ofile name")
    parser.add_argument("--by","-b",default="site",help="add tag by gene or site,default is site")
    args = vars(parser.parse_args())
    init_file_list = args.get("init")
    add_tag_file = args.get("add_tag")
    infile = args.get("infile")
    ofile = args.get("ofile")
    by = args.get("by")
    if init_file_list and add_tag_file:
        raise Exception("--init and --add_tag can not present at the same time!")
    if init_file_list:
        init_index(init_file_list,ofile,by)
    elif add_tag_file:
        add_tags(infile,ofile,add_tag_file,by)
    else:
        raise Exception("--init or --add_tag is required")
